{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reduction Operators"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This section shows you how to use operators from `reductions` module. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. Start with a simple pipeline based on `ExternalSource`. Input has two samples per batch. Shape of both samples is (3, 3). First contains consecutive numbers, second contains consecutive even numbers. This will be useful to visualize possible reductions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output 0:\n",
      "[[[ 0  1  2]\n",
      "  [ 3  4  5]\n",
      "  [ 6  7  8]]\n",
      "\n",
      " [[ 0  2  4]\n",
      "  [ 6  8 10]\n",
      "  [12 14 16]]] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "import nvidia.dali.fn as fn\n",
    "import nvidia.dali.types as types\n",
    "import nvidia.dali.backend as backend\n",
    "from nvidia.dali.pipeline import Pipeline\n",
    "import numpy as np\n",
    "\n",
    "batch_size = 2\n",
    "\n",
    "\n",
    "def get_batch():\n",
    "    return [\n",
    "        np.reshape(np.arange(9), (3, 3)) * (i + 1) for i in range(batch_size)\n",
    "    ]\n",
    "\n",
    "\n",
    "def run_and_print(pipe):\n",
    "    pipe.build()\n",
    "    output = pipe.run()\n",
    "    for i, out in enumerate(output):\n",
    "        if type(out) == backend.TensorListGPU:\n",
    "            out = out.as_cpu()\n",
    "        output_array = out.as_array()\n",
    "        print(\"Output {}:\\n{} \\n\".format(i, output_array))\n",
    "\n",
    "\n",
    "pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)\n",
    "with pipe:\n",
    "    input = fn.external_source(source=get_batch, dtype=types.INT64)\n",
    "\n",
    "    pipe.set_outputs(input)\n",
    "\n",
    "run_and_print(pipe)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. Add some reductions to the pipeline above. Begin with the `Max` operator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output 0:\n",
      "[ 8 16] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)\n",
    "with pipe:\n",
    "    input = fn.external_source(source=get_batch, dtype=types.INT64)\n",
    "    max = fn.reductions.max(input)\n",
    "\n",
    "    pipe.set_outputs(max)\n",
    "\n",
    "run_and_print(pipe)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, it returned the biggest value from each sample. \n",
    "\n",
    "3. Perform other reductions like `Min` or `Sum`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output 0:\n",
      "[0 0] \n",
      "\n",
      "Output 1:\n",
      "[36 72] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)\n",
    "with pipe:\n",
    "    input = fn.external_source(source=get_batch, dtype=types.INT64)\n",
    "    min = fn.reductions.min(input)\n",
    "    sum = fn.reductions.sum(input)\n",
    "\n",
    "    pipe.set_outputs(min, sum)\n",
    "\n",
    "run_and_print(pipe)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the code samples above we see reductions performed for all elements of each sample. \n",
    "\n",
    "4. Reductions can be performed along an arbitrary set of axes. To control this behavior you can use `axes` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output 0:\n",
      "[[0 1 2]\n",
      " [0 2 4]] \n",
      "\n",
      "Output 1:\n",
      "[[ 0  3  6]\n",
      " [ 0  6 12]] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)\n",
    "with pipe:\n",
    "    input = fn.external_source(source=get_batch, dtype=types.INT64)\n",
    "    min_axis_0 = fn.reductions.min(input, axes=0)\n",
    "    min_axis_1 = fn.reductions.min(input, axes=1)\n",
    "\n",
    "    pipe.set_outputs(min_axis_0, min_axis_1)\n",
    "\n",
    "run_and_print(pipe)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Min` reduction was performed along axis 0 and 1, and it returned minimum element per column and per row respectively. \n",
    "\n",
    "To make it easier, reductions support `axis_names` argument. It allows to pass axis names rather than indices. Names are matched based on the layout of the input. You need to provide layout argument in `ExternalSource`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output 0:\n",
      "[[0 1 2]\n",
      " [0 2 4]] \n",
      "\n",
      "Output 1:\n",
      "[[ 0  3  6]\n",
      " [ 0  6 12]] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)\n",
    "with pipe:\n",
    "    input = fn.external_source(source=get_batch, layout=\"AB\", dtype=types.INT64)\n",
    "    min_axis_0 = fn.reductions.min(input, axis_names=\"A\")\n",
    "    min_axis_1 = fn.reductions.min(input, axis_names=\"B\")\n",
    "\n",
    "    pipe.set_outputs(min_axis_0, min_axis_1)\n",
    "\n",
    "run_and_print(pipe)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note**: Passing all axes will result in a full reduction, while passing empty axes will result in no reduction. This is true for both indices and layouts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output 0:\n",
      "[0 0] \n",
      "\n",
      "Output 1:\n",
      "[[[ 0  1  2]\n",
      "  [ 3  4  5]\n",
      "  [ 6  7  8]]\n",
      "\n",
      " [[ 0  2  4]\n",
      "  [ 6  8 10]\n",
      "  [12 14 16]]] \n",
      "\n",
      "Output 2:\n",
      "[0 0] \n",
      "\n",
      "Output 3:\n",
      "[[[ 0  1  2]\n",
      "  [ 3  4  5]\n",
      "  [ 6  7  8]]\n",
      "\n",
      " [[ 0  2  4]\n",
      "  [ 6  8 10]\n",
      "  [12 14 16]]] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)\n",
    "with pipe:\n",
    "    input = fn.external_source(source=get_batch, layout=\"AB\", dtype=types.INT64)\n",
    "    min_axes_full = fn.reductions.min(input, axes=(0, 1))\n",
    "    min_axes_empty = fn.reductions.min(input, axes=())\n",
    "    min_layout_full = fn.reductions.min(input, axis_names=\"AB\")\n",
    "    min_layout_empty = fn.reductions.min(input, axis_names=\"\")\n",
    "\n",
    "    pipe.set_outputs(\n",
    "        min_axes_full, min_axes_empty, min_layout_full, min_layout_empty\n",
    "    )\n",
    "\n",
    "run_and_print(pipe)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5. For inputs with higher dimensionality you can pass any combination of the axes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_batch():\n",
    "    return [\n",
    "        np.reshape(np.arange(8, dtype=np.int32), (2, 2, 2)) * (i + 1)\n",
    "        for i in range(batch_size)\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output 0:\n",
      "[[[[ 0  1]\n",
      "   [ 2  3]]\n",
      "\n",
      "  [[ 4  5]\n",
      "   [ 6  7]]]\n",
      "\n",
      "\n",
      " [[[ 0  2]\n",
      "   [ 4  6]]\n",
      "\n",
      "  [[ 8 10]\n",
      "   [12 14]]]] \n",
      "\n",
      "Output 1:\n",
      "[[0 1]\n",
      " [0 2]] \n",
      "\n",
      "Output 2:\n",
      "[[0 2]\n",
      " [0 4]] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)\n",
    "with pipe:\n",
    "    input = fn.external_source(\n",
    "        source=get_batch, layout=\"ABC\", dtype=types.INT32\n",
    "    )\n",
    "    min_axes_empty = fn.reductions.min(input, axes=())\n",
    "    min_axes_0_1 = fn.reductions.min(input, axes=(0, 1))\n",
    "    min_layout_A_C = fn.reductions.min(input, axis_names=\"AC\")\n",
    "\n",
    "    pipe.set_outputs(min_axes_empty, min_axes_0_1, min_layout_A_C)\n",
    "\n",
    "run_and_print(pipe)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6. There are reductions that require additional inputs. `StdDev` and `Variance` rely on the externally provided mean, and it can be calculated with the `Mean` reduction operator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output 0:\n",
      "[3.5 7. ] \n",
      "\n",
      "Output 1:\n",
      "[2.291288 4.582576] \n",
      "\n",
      "Output 2:\n",
      "[ 5.25 21.  ] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)\n",
    "with pipe:\n",
    "    input = fn.external_source(source=get_batch, dtype=types.INT32)\n",
    "    mean = fn.reductions.mean(input)\n",
    "    std_dev = fn.reductions.std_dev(input, mean)\n",
    "    variance = fn.reductions.variance(input, mean)\n",
    "\n",
    "    pipe.set_outputs(mean, std_dev, variance)\n",
    "\n",
    "run_and_print(pipe)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "7. By default, reductions remove unnecessary dimensions. This behaviour can be controlled with the `keep_dims` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output 0:\n",
      "[3.5 7. ] \n",
      "\n",
      "Output 1:\n",
      "[2.291288 4.582576] \n",
      "\n",
      "Output 2:\n",
      "[ 5.25 21.  ] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)\n",
    "with pipe:\n",
    "    input = fn.external_source(source=get_batch, dtype=types.INT32)\n",
    "    mean = fn.reductions.mean(input)\n",
    "    std_dev = fn.reductions.std_dev(input, mean, keep_dims=True)\n",
    "    variance = fn.reductions.variance(input, mean)\n",
    "\n",
    "    pipe.set_outputs(mean, std_dev, variance)\n",
    "\n",
    "run_and_print(pipe)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the code sample above applying reductions resulted in changing the output type. \n",
    "\n",
    "8. The argument `dtype` can be used to specify the desired output data type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output 0:\n",
      "[28 56] \n",
      "\n",
      "Output 1:\n",
      "[28. 56.] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)\n",
    "with pipe:\n",
    "    input = fn.external_source(source=get_batch, dtype=types.INT32)\n",
    "    sum_int_64 = fn.reductions.sum(input, dtype=types.INT64)\n",
    "    sum_float = fn.reductions.sum(input, dtype=types.FLOAT)\n",
    "\n",
    "    pipe.set_outputs(sum_int_64, sum_float)\n",
    "\n",
    "run_and_print(pipe)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note**: Not all data types combinations are supported. The default behaviour varies from operator to operator. A general rule is for the output type to be able to accommodate the result depending on the input type. For example, for the input type `INT32` the default output type of a sum is `INT32` and the default output type of a mean is `FLOAT`.\n",
    "\n",
    "9. All reductions can be offloaded to the GPU. GPU variants work the same way as their CPU counterparts. Below we show a code sample containing all reductions offloaded to the GPU with various parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output 0:\n",
      "[[[[0]\n",
      "   [2]]]\n",
      "\n",
      "\n",
      " [[[0]\n",
      "   [4]]]] \n",
      "\n",
      "Output 1:\n",
      "[[[[ 7]]]\n",
      "\n",
      "\n",
      " [[[14]]]] \n",
      "\n",
      "Output 2:\n",
      "[28 56] \n",
      "\n",
      "Output 3:\n",
      "[[[ 2.  3.]\n",
      "  [ 4.  5.]]\n",
      "\n",
      " [[ 4.  6.]\n",
      "  [ 8. 10.]]] \n",
      "\n",
      "Output 4:\n",
      "[17.5 70. ] \n",
      "\n",
      "Output 5:\n",
      "[[[[ 0.  1.]\n",
      "   [ 2.  3.]]\n",
      "\n",
      "  [[ 4.  5.]\n",
      "   [ 6.  7.]]]\n",
      "\n",
      "\n",
      " [[[ 0.  2.]\n",
      "   [ 4.  6.]]\n",
      "\n",
      "  [[ 8. 10.]\n",
      "   [12. 14.]]]] \n",
      "\n",
      "Output 6:\n",
      "[[[2. 2.]\n",
      "  [2. 2.]]\n",
      "\n",
      " [[4. 4.]\n",
      "  [4. 4.]]] \n",
      "\n",
      "Output 7:\n",
      "[[[[ 4.  4.]\n",
      "   [ 4.  4.]]]\n",
      "\n",
      "\n",
      " [[[16. 16.]\n",
      "   [16. 16.]]]] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)\n",
    "with pipe:\n",
    "    input = fn.external_source(\n",
    "        source=get_batch, layout=\"ABC\", dtype=types.INT32\n",
    "    )\n",
    "    min = fn.reductions.min(input.gpu(), axis_names=\"AC\", keep_dims=True)\n",
    "    max = fn.reductions.max(input.gpu(), keep_dims=True)\n",
    "    sum = fn.reductions.sum(input.gpu(), dtype=types.INT64)\n",
    "    mean = fn.reductions.mean(input.gpu(), axes=0)\n",
    "    mean_square = fn.reductions.mean_square(input.gpu())\n",
    "    rms = fn.reductions.rms(input.gpu(), axes=(), dtype=types.FLOAT)\n",
    "    std_dev = fn.reductions.std_dev(input.gpu(), mean, axes=0)\n",
    "    variance = fn.reductions.variance(\n",
    "        input.gpu(), mean.gpu(), axes=0, keep_dims=True\n",
    "    )\n",
    "\n",
    "    pipe.set_outputs(min, max, sum, mean, mean_square, rms, std_dev, variance)\n",
    "\n",
    "run_and_print(pipe)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
